import os
import sqlite3
import hashlib
from pathlib import Path
from typing import Optional, Dict


class DatabaseMerger:
    def __init__(self, merged_db="msg/merged.db", contact_db="contact.db"):
        self.merged_db = merged_db
        self.contact_db = contact_db
        self.batch_size = 10000  # 根据内存调整批次大小

        # 确保msg目录存在
        msg_dir = os.path.dirname(merged_db)
        if msg_dir and not os.path.exists(msg_dir):
            os.makedirs(msg_dir)

        # 初始化目标数据库
        with self._connect(self.merged_db) as conn:
            conn.execute('''
                         CREATE TABLE IF NOT EXISTS merged_msg
                         (
                             source_table
                             TEXT,
                             local_id
                             INTEGER,
                             server_id
                             INTEGER,
                             local_type
                             INTEGER,
                             sort_seq
                             INTEGER,
                             real_sender_id
                             INTEGER,
                             create_time
                             INTEGER,
                             status
                             INTEGER,
                             upload_status
                             INTEGER,
                             download_status
                             INTEGER,
                             server_seq
                             INTEGER,
                             origin_source
                             INTEGER,
                             source
                             TEXT,
                             message_content
                             TEXT,
                             compress_content
                             TEXT,
                             packed_info_data
                             BLOB,
                             WCDB_CT_message_content
                             INTEGER,
                             WCDB_CT_source
                             INTEGER,
                             sender_name
                             TEXT,
                             source_md5
                             TEXT
                             GENERATED
                             ALWAYS AS (
                             CASE
                             WHEN
                             instr
                         (
                             source_table,
                             '_'
                         ) > 0
                             THEN substr
                         (
                             source_table,
                             instr
                         (
                             source_table,
                             '_'
                         )+1)
                             ELSE NULL END
                             ) VIRTUAL,
                             username TEXT,
                             remark TEXT,
                             nick_name TEXT
                             )
                         ''')
            conn.execute("CREATE INDEX IF NOT EXISTS idx_md5 ON merged_msg(source_md5)")
            conn.commit()

    def _connect(self, db_path, optimze=False):
        """创建数据库连接"""
        conn = sqlite3.connect(db_path)
        if optimze:
            conn.execute("PRAGMA synchronous = OFF")
            conn.execute("PRAGMA journal_mode = MEMORY")
            conn.execute("PRAGMA cache_size = -10000")  # 10MB缓存
        return conn

    def _get_name_mapping(self, conn):
        """获取Name2Id映射表"""
        try:
            return dict(conn.execute("SELECT rowid, user_name FROM Name2Id").fetchall())
        except sqlite3.OperationalError:
            return {}

    def merge_sources(self, source_dirs, table_pattern="Msg_%"):
        """合并源数据库"""
        global current_progress
        current_progress = 0

        # 收集所有源数据库
        source_dbs = []
        for dir_path in source_dirs:
            # 查找所有可能的消息数据库
            source_dbs.extend(Path(dir_path).rglob("message_*.db"))
            source_dbs.extend(Path(dir_path).rglob("MSG*.db"))
            source_dbs.extend(Path(dir_path).rglob("ChatMsg.db"))

        total_dbs = len(source_dbs)
        if total_dbs == 0:
            print("未找到任何消息数据库文件")
            return False

        progress_per_db = 70 / total_dbs  # 数据库合并占70%的进度

        # 分批处理每个数据库
        with self._connect(self.merged_db, optimze=True) as dest_conn:
            dest_cur = dest_conn.cursor()

            for db_index, db_path in enumerate(source_dbs):
                if not db_path.exists():
                    continue

                print(f"处理数据库: {db_path}")
                with self._connect(db_path) as src_conn:
                    # 获取姓名映射
                    name_map = self._get_name_mapping(src_conn)

                    # 获取所有表
                    tables = [
                        t[0] for t in
                        src_conn.execute(
                            "SELECT name FROM sqlite_master WHERE type='table'"
                        )
                    ]

                    msg_tables = [t for t in tables if t.startswith('Msg_')]
                    # 如果没有Msg_表，则处理所有可能包含消息的表
                    if not msg_tables:
                        # 添加其他可能的消息表
                        msg_tables = [t for t in tables if
                                      any(keyword in t.lower() for keyword in ['msg', 'message', 'chat'])]

                    print(f"找到可能的消息表: {msg_tables}")

                    for table_index, table in enumerate(msg_tables):
                        try:
                            # 首先检查表结构
                            columns_info = src_conn.execute(f"PRAGMA table_info({table})").fetchall()
                            column_names = [col[1] for col in columns_info]
                            column_map = {col[1]: col[0] for col in columns_info}  # name -> index
                            print(f"表 {table} 的列: {column_names}")

                            # 定义列映射关系（处理不同版本的数据库）
                            column_mapping = {
                                'local_id': ['local_id', 'localId', 'msgLocalId', 'MsgSvrID'],
                                'server_id': ['server_id', 'MsgSvrID'],
                                'local_type': ['local_type', 'Type'],
                                'create_time': ['create_time', 'CreateTime'],
                                'status': ['status', 'Status'],
                                'source': ['source', 'StrTalker', 'Talker'],
                                'message_content': ['message_content', 'StrContent', 'Content'],
                                'real_sender_id': ['real_sender_id', 'TalkerId', 'talkerId'],
                                'sort_seq': ['sort_seq', 'Sequence'],
                                'compress_content': ['compress_content', 'CompressContent'],
                                'packed_info_data': ['packed_info_data', 'BytesExtra'],
                            }

                            # 构建实际可用的列映射
                            actual_mapping: Dict[str, str] = {}
                            for standard_col, possible_cols in column_mapping.items():
                                for col in possible_cols:
                                    if col in column_names:
                                        actual_mapping[standard_col] = col
                                        break

                            print(f"表 {table} 的列映射: {actual_mapping}")

                            # 检查是否有必要字段
                            required_fields = ['local_id', 'create_time']
                            if not all(field in actual_mapping for field in required_fields):
                                missing = [f for f in required_fields if f not in actual_mapping]
                                print(f"跳过表 {table}，缺少关键字段: {missing}")
                                continue

                            # 构建SELECT语句
                            select_columns = [actual_mapping[field] for field in actual_mapping.keys()]
                            select_clause = ', '.join(select_columns)

                            src_cur = src_conn.cursor()
                            src_cur.execute(f'SELECT {select_clause} FROM {table}')

                            insert_buffer = []
                            while True:
                                rows = src_cur.fetchmany(self.batch_size)
                                if not rows:
                                    break

                                # 处理行数据，映射到标准格式
                                for row in rows:
                                    row_dict = dict(zip(select_columns, row))

                                    # 构建标准格式的行
                                    standard_row = [
                                        row_dict.get(actual_mapping.get('local_id')),
                                        row_dict.get(actual_mapping.get('server_id')),
                                        row_dict.get(actual_mapping.get('local_type')),
                                        row_dict.get(actual_mapping.get('sort_seq')),
                                        row_dict.get(actual_mapping.get('real_sender_id')),
                                        row_dict.get(actual_mapping.get('create_time')),
                                        row_dict.get(actual_mapping.get('status')),
                                        None,  # upload_status
                                        None,  # download_status
                                        None,  # server_seq
                                        None,  # origin_source
                                        row_dict.get(actual_mapping.get('source')),
                                        row_dict.get(actual_mapping.get('message_content')),
                                        row_dict.get(actual_mapping.get('compress_content')),
                                        row_dict.get(actual_mapping.get('packed_info_data')),
                                        None,  # WCDB_CT_message_content
                                        None,  # WCDB_CT_source
                                    ]

                                    # 添加到插入缓冲区
                                    insert_buffer.append((
                                        table,  # source_table
                                        *standard_row,
                                        None,  # sender_name (暂时设为None)
                                    ))

                                # 批量插入
                                if len(insert_buffer) >= self.batch_size:
                                    dest_cur.executemany('''
                                                         INSERT INTO merged_msg (source_table, local_id, server_id,
                                                                                 local_type,
                                                                                 sort_seq, real_sender_id, create_time,
                                                                                 status,
                                                                                 upload_status, download_status,
                                                                                 server_seq,
                                                                                 origin_source, source, message_content,
                                                                                 compress_content, packed_info_data,
                                                                                 WCDB_CT_message_content,
                                                                                 WCDB_CT_source,
                                                                                 sender_name)
                                                         VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
                                                                 ?)
                                                         ''', insert_buffer)
                                    dest_conn.commit()
                                    insert_buffer.clear()

                            # 插入剩余数据
                            if insert_buffer:
                                dest_cur.executemany('''
                                                     INSERT INTO merged_msg (source_table, local_id, server_id,
                                                                             local_type,
                                                                             sort_seq, real_sender_id, create_time,
                                                                             status,
                                                                             upload_status, download_status, server_seq,
                                                                             origin_source, source, message_content,
                                                                             compress_content, packed_info_data,
                                                                             WCDB_CT_message_content, WCDB_CT_source,
                                                                             sender_name)
                                                     VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                                                     ''', insert_buffer)
                                dest_conn.commit()

                        except sqlite3.Error as e:
                            print(f"处理表 {table} 时出错: {str(e)}")

                # 更新进度
                current_progress = int((db_index + 1) * progress_per_db)

        return True

    def update_contacts(self):
        """更新联系人信息"""
        global current_progress

        # 预加载联系人数据
        contact_map = {}

        # 查找所有可能的联系人数据库
        all_db_files = list(Path("./db").rglob("*.db"))
        print(f"找到所有数据库文件: {[str(f) for f in all_db_files]}")

        # 尝试不同的联系人表和列名组合
        contact_table_patterns = [
            ("Contact", ["UserName", "Remark", "NickName"]),
            ("Contact", ["username", "remark", "nick_name"]),
            ("Contact", ["username", "remark", "nickname"]),
            ("Friend", ["UserName", "Remark", "NickName"]),
            ("UserInfo", ["UserName", "Remark", "NickName"]),
        ]

        for db_file in all_db_files:
            print(f"尝试从 {db_file} 读取联系人信息")
            try:
                with self._connect(db_file) as conn:
                    # 检查数据库中有哪些表
                    tables = conn.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()
                    table_names = [t[0] for t in tables]
                    print(f"  {db_file} 中的表: {table_names}")

                    # 跳过系统表
                    user_tables = [t for t in table_names if not t.startswith('sqlite_')]
                    if not user_tables:
                        continue

                    for table_name, column_names in contact_table_patterns:
                        if table_name not in table_names:
                            continue

                        try:
                            # 检查表是否有数据
                            count_cursor = conn.execute(f"SELECT COUNT(*) FROM {table_name}")
                            count = count_cursor.fetchone()[0]
                            if count == 0:
                                continue

                            print(f"  在 {db_file} 中找到 {table_name} 表，包含 {count} 条记录")

                            # 尝试读取数据
                            username_col, remark_col, nickname_col = column_names
                            query = f"SELECT {username_col}, {remark_col}, {nickname_col} FROM {table_name} WHERE {username_col} IS NOT NULL"
                            rows = conn.execute(query).fetchall()

                            added_count = 0
                            for row in rows:
                                username, remark, nick_name = row
                                if username:
                                    md5 = hashlib.md5(username.encode()).hexdigest()
                                    if md5 not in contact_map:
                                        contact_map[md5] = (username, remark or "", nick_name or "")
                                        added_count += 1

                            if added_count > 0:
                                print(f"    从 {db_file} 的 {table_name} 表中成功加载 {added_count} 个新联系人")

                        except sqlite3.OperationalError as e:
                            print(f"    尝试 {table_name} 表时出错: {str(e)}")
                            continue
            except Exception as e:
                print(f"读取 {db_file} 时出错: {str(e)}")
                continue

        if not contact_map:
            print("未在任何数据库中找到有效的联系人表")
            return False

        print(f"成功加载 {len(contact_map)} 个联系人")

        # 分批更新
        with self._connect(self.merged_db, optimze=True) as conn:
            conn.execute("BEGIN")
            try:
                cursor = conn.cursor()
                cursor.execute("SELECT rowid, source_md5 FROM merged_msg")

                total_rows = cursor.fetchall()
                if not total_rows:
                    current_progress = 100
                    return True

                progress_per_batch = 30 / (len(total_rows) / self.batch_size) if len(
                    total_rows) > self.batch_size else 30  # 联系人更新占30%的进度

                update_buffer = []
                for batch_index, (rowid, md5) in enumerate(total_rows):
                    if md5 in contact_map:
                        update_buffer.append((*contact_map[md5], rowid))

                    if len(update_buffer) >= self.batch_size:
                        conn.executemany('''
                                         UPDATE merged_msg
                                         SET username=?,
                                             remark=?,
                                             nick_name=?
                                         WHERE rowid = ?
                                         ''', update_buffer)
                        update_buffer.clear()
                        # 更新进度
                        current_progress = 70 + int((batch_index / len(total_rows)) * 30)

                if update_buffer:
                    conn.executemany('''
                                     UPDATE merged_msg
                                     SET username=?,
                                         remark=?,
                                         nick_name=?
                                     WHERE rowid = ?
                                     ''', update_buffer)

                conn.commit()
                current_progress = 100
            except Exception as e:
                conn.rollback()
                print(f"更新联系人信息时出错: {str(e)}")
                raise e

        return True


def find(directory, glob_pattern):
    """查找文件"""
    root_dir = Path(directory)
    return [item for item in root_dir.rglob(glob_pattern) if item.is_file()]


def main():
    """主处理函数"""
    try:
        # 在整个db目录中查找所有数据库文件
        all_dbs = find("./db", "*.db")
        contact_db_fp = [db for db in all_dbs if 'contact' in str(db).lower()]
        message_dbs = [db for db in all_dbs if
                       any(keyword in str(db).lower() for keyword in ['message', 'msg', 'chat'])]

        # 如果没有明确的消息数据库，就把所有数据库都当作候选
        if not message_dbs:
            message_dbs = all_dbs

        if all_dbs:
            print(f"找到数据库文件 {len(all_dbs)} 个")
            print(f"联系人相关数据库: {len(contact_db_fp)} 个")
            print(f"消息相关数据库: {len(message_dbs)} 个")

            merger = DatabaseMerger('msg/merged.db')

            # Step 1: 合并数据
            if not merger.merge_sources(source_dirs=['./db']):
                return False

            # Step 2: 更新联系人信息
            if not merger.update_contacts():
                return False

            return True
        else:
            print("未找到任何数据库文件")
            return False

    except Exception as e:
        print(f"处理过程中发生错误: {str(e)}")
        return False